from datetime import datetime

import pandas as pd
import numpy as np
import xarray as xr

import ttide
from ttide.t_tide import t_tide

from ppgnss import gnss_time
from ppgnss import gnss_utils

from plot_baseline import read_seal_file, read_gamit_solution
import matplotlib.pyplot as plt

gamit_solution = "../baseline/solution.dat"

excel_file_path = "../data/cczihe.xlsx"
seal_down_file = "../data/sea_levels/downstream_7308.txt"
seal_up_file = "../data/sea_levels/upstream_3020.txt"


df = pd.read_excel(excel_file_path)
xr_gg = read_gamit_solution(gamit_solution, 2022, 118)
xr_seal_down = read_seal_file(seal_down_file)
xr_seal_up = read_seal_file(seal_up_file)

# ['采集时间', '设备序列号', '东向坐标 (m)', '北向坐标 (m)', '天向坐标 (m)', 
# '东向坐标日变形速率 (mm/d)',  '北向坐标日变形速率 (mm/d)', '天向坐标日变形速率 (mm/d)',
# '平面日变形速率 (mm/d)', '三维日变形速率 (mm/d)', '差分期龄 (s)']
count = len(df["采集时间"])
time_from = np.datetime64(df["采集时间"][count-1])
time_to = np.datetime64(df["采集时间"][0])
dtime = np.datetime64(df["采集时间"][count-2])-np.datetime64(df["采集时间"][count-1])
nepoch = (time_to - time_from)//dtime+1

times = np.array([time_from + dtime*_ - np.timedelta64(8, "h") for _ in range(nepoch)])

dus = df["天向坐标 (m)"]-df["天向坐标 (m)"][0]
dns = df["北向坐标 (m)"]-df["北向坐标 (m)"][0]
des = df["东向坐标 (m)"]-df["东向坐标 (m)"][0]
idx = np.logical_or(np.logical_or(np.abs(dus) > 0.04, np.abs(dns) > 0.04), np.abs(des)>0.04)
dus[idx] = np.nan
dns[idx] = np.nan
des[idx] = np.nan

shape = (len(times), 3)
data = np.full(shape, np.nan)

xr_data = xr.DataArray(data, dims=["time", "data"], coords=[times, ['n', 'e', 'u']])
valid_times = np.array([np.datetime64(_) - np.timedelta64(8, "h") for _ in df["采集时间"]]) 

xr_data.loc[valid_times, "u"] = dus
xr_data.loc[valid_times, "n"] = dns
xr_data.loc[valid_times, "e"] = des

# xr_time1, xr_time2 = xr.align(xr_data.coords["time"], xr_seal_down.coords["time"])
# print(xr_time1)
time_resample = [np.datetime64(str(times[0])[:10]+" 00:00:00")+np.timedelta64(_, "h") for _ in range(500*24)]

xr_data, xr_seal_down = xr.align(xr_data, xr_seal_down, join="left")
xr_data, xr_seal_up = xr.align(xr_data, xr_seal_up, join="left")
xr_seal_delta = xr_seal_up - xr_seal_down
# xr_data = xr_data.resample(time='12H').mean()

# timelist = [_ for _ in time_resample if _ in xr_data.coords["time"].values]
# xr_data = xr_data.loc[timelist]
# xr_seal = xr_seal.loc[timelist]
time_from = "2022-01-01"
time_to = "2023-03-30"
# plt.plot(xr_data.loc[time_from:time_to].coords["time"], xr_data.loc[time_from:time_to, "u"], label="U")
plt.plot(xr_data.loc[time_from:time_to].coords["time"], xr_data.loc[time_from:time_to, "n"], label="N")
plt.plot(xr_data.loc[time_from:time_to].coords["time"], xr_data.loc[time_from:time_to, "e"], label="E")
# plt.plot(xr_gg.loc[time_from:time_to].coords["time"], xr_gg.loc[time_from:time_to, "e"], label="E")
# plt.plot(xr_gg.loc[time_from:time_to].coords["time"], xr_gg.loc[time_from:time_to, "n"], label="N")

plt.legend()

ax = plt.twinx()
plt.plot(xr_seal_down.loc[time_from:time_to].coords["time"], xr_seal_down.loc[time_from:time_to], label="SL", c="r", alpha=0.3)
# plt.plot(xr_seal_delta.loc[time_from:time_to].coords["time"], xr_seal_delta.loc[time_from:time_to], label="SL", c="r", alpha=0.3)
# plt.plot(xr_seal_up.loc[time_from:time_to].coords["time"], xr_seal_up.loc[time_from:time_to], label="SL", c="r", alpha=0.3)

plt.grid(True)
# plt.show()
plt.close()
for com in ["u", "e", "n"]:
    stime = datetime.strptime(str(xr_data.coords["time"].values[0])[:19], '%Y-%m-%dT%H:%M:%S')
    obs1 = xr_data.loc[time_from:time_to, com]
    # obs1 = obs1 - obs1.mean()
    obs1 = obs1.interpolate_na(dim="time")

    out1 = t_tide(obs1, dt=1, stime=stime, secular="linear", lat=np.array(30.23))
    print(np.std(obs1), np.std(out1["xout"]))
    fig, ax = plt.subplots(nrows=1, sharey=True, sharex=True, figsize=(13, 5))
    ax0, ax1, ax2 = ax, ax, ax
    ax0.plot(obs1.coords["time"], obs1, label=u'Observations')
    ax0.legend(numpoints=1, loc='lower right')

    ax1.plot(obs1.coords["time"], out1["xout"].squeeze(), alpha=0.5, label=u'Prediction')
    ax1.legend(numpoints=1, loc='lower right')

    # ax2.plot(obs.coords["time"], obs-out["xout"].squeeze(), alpha=0.5, label=u'Residue')
    # _ = ax2.legend(numpoints=1, loc='lower right')
    plt.title(com)
    plt.show()

xr_seal_down = xr_seal_down.loc[time_from:time_to]
obs = xr_seal_down - xr_seal_down.mean()
obs = obs.interpolate_na(dim="time")
print(str(obs.coords["time"].values[0]))
stime = datetime.strptime(str(obs.coords["time"].values[0])[:19], '%Y-%m-%dT%H:%M:%S')
out = t_tide(obs, dt=1, stime=stime, lat=np.array(30.23))
# print(tTideCon)

fig, (ax0, ax1, ax2) = plt.subplots(nrows=3, sharey=True, sharex=True, figsize=(13, 5))

ax0.plot(obs.coords["time"], obs, label=u'Observations')
ax0.legend(numpoints=1, loc='lower right')

ax1.plot(obs.coords["time"], out["xout"].squeeze(), alpha=0.5, label=u'Prediction')
ax1.legend(numpoints=1, loc='lower right')

ax1.plot(obs1.coords["time"], out1["xout"].squeeze(), alpha=0.5, label=u'N')

# ax2.plot(obs.coords["time"], obs-out["xout"].squeeze(), alpha=0.5, label=u'Residue')
# _ = ax2.legend(numpoints=1, loc='lower right')
plt.show()
